"""
数据库迁移运行器
"""

import asyncio
from typing import Dict, Any, List, Optional
from datetime import datetime
import logging

from ..core.database import get_db_session, Base, engine
from ..core.logging import get_logger

logger = get_logger(__name__)


class MigrationRunner:
    """迁移运行器"""
    
    def __init__(self):
        self.migrations: Dict[str, Any] = {}
        self.applied_migrations: List[str] = []
    
    def register_migration(self, version: str, migration_func: Any):
        """注册迁移"""
        self.migrations[version] = migration_func
    
    async def init_migration_table(self):
        """初始化迁移表"""
        try:
            async with get_db_session() as session:
                # 创建迁移记录表
                await session.execute("""
                    CREATE TABLE IF NOT EXISTS migration_history (
                        version VARCHAR(50) PRIMARY KEY,
                        applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        description TEXT
                    )
                """)
                await session.commit()
                
                # 获取已应用的迁移
                result = await session.execute("SELECT version FROM migration_history ORDER BY version")
                rows = result.fetchall()
                self.applied_migrations = [row[0] for row in rows]
                
        except Exception as e:
            logger.error(f"Failed to initialize migration table: {e}")
            raise
    
    async def get_pending_migrations(self) -> List[str]:
        """获取待执行的迁移"""
        all_versions = sorted(self.migrations.keys())
        pending = [v for v in all_versions if v not in self.applied_migrations]
        return pending
    
    async def apply_migration(self, version: str):
        """应用指定版本的迁移"""
        if version not in self.migrations:
            raise ValueError(f"Migration {version} not found")
        
        if version in self.applied_migrations:
            logger.info(f"Migration {version} already applied")
            return
        
        migration = self.migrations[version]
        
        try:
            logger.info(f"Applying migration {version}")
            
            # 执行迁移
            if hasattr(migration, 'up'):
                if asyncio.iscoroutinefunction(migration.up):
                    await migration.up()
                else:
                    migration.up()
            elif callable(migration):
                if asyncio.iscoroutinefunction(migration):
                    await migration()
                else:
                    migration()
            
            # 记录迁移
            async with get_db_session() as session:
                description = getattr(migration, 'description', f"Migration {version}")
                await session.execute(
                    "INSERT INTO migration_history (version, description) VALUES (?, ?)",
                    (version, description)
                )
                await session.commit()
            
            self.applied_migrations.append(version)
            logger.info(f"Migration {version} applied successfully")
            
        except Exception as e:
            logger.error(f"Failed to apply migration {version}: {e}")
            raise
    
    async def rollback_migration(self, version: str):
        """回滚指定版本的迁移"""
        if version not in self.applied_migrations:
            logger.warning(f"Migration {version} not applied, cannot rollback")
            return
        
        if version not in self.migrations:
            raise ValueError(f"Migration {version} not found")
        
        migration = self.migrations[version]
        
        try:
            logger.info(f"Rolling back migration {version}")
            
            # 执行回滚
            if hasattr(migration, 'down'):
                if asyncio.iscoroutinefunction(migration.down):
                    await migration.down()
                else:
                    migration.down()
            
            # 删除迁移记录
            async with get_db_session() as session:
                await session.execute(
                    "DELETE FROM migration_history WHERE version = ?",
                    (version,)
                )
                await session.commit()
            
            self.applied_migrations.remove(version)
            logger.info(f"Migration {version} rolled back successfully")
            
        except Exception as e:
            logger.error(f"Failed to rollback migration {version}: {e}")
            raise
    
    async def migrate(self, target_version: Optional[str] = None):
        """执行迁移到指定版本"""
        await self.init_migration_table()
        
        if target_version:
            # 迁移到指定版本
            if target_version not in self.migrations:
                raise ValueError(f"Target version {target_version} not found")
            
            all_versions = sorted(self.migrations.keys())
            target_index = all_versions.index(target_version)
            
            if target_index < len(self.applied_migrations):
                # 需要回滚
                for i in range(len(self.applied_migrations) - 1, target_index - 1, -1):
                    version = self.applied_migrations[i]
                    await self.rollback_migration(version)
            else:
                # 需要前进
                for version in all_versions[len(self.applied_migrations):target_index + 1]:
                    await self.apply_migration(version)
        else:
            # 迁移到最新版本
            pending = await self.get_pending_migrations()
            for version in pending:
                await self.apply_migration(version)
    
    async def rollback(self, target_version: str):
        """回滚到指定版本"""
        await self.init_migration_table()
        
        if target_version not in self.migrations:
            raise ValueError(f"Target version {target_version} not found")
        
        # 获取需要回滚的版本（倒序）
        to_rollback = []
        for version in reversed(self.applied_migrations):
            to_rollback.append(version)
            if version == target_version:
                break
        
        if target_version not in self.applied_migrations:
            raise ValueError(f"Version {target_version} not applied")
        
        # 执行回滚
        for version in to_rollback:
            if version != target_version:
                await self.rollback_migration(version)
    
    async def get_status(self) -> Dict[str, Any]:
        """获取迁移状态"""
        await self.init_migration_table()
        
        all_versions = sorted(self.migrations.keys())
        applied_versions = sorted(self.applied_migrations)
        pending_versions = await self.get_pending_migrations()
        
        return {
            "total_migrations": len(all_versions),
            "applied_migrations": len(applied_versions),
            "pending_migrations": len(pending_versions),
            "all_versions": all_versions,
            "applied_versions": applied_versions,
            "pending_versions": pending_versions,
            "current_version": applied_versions[-1] if applied_versions else None,
            "latest_version": all_versions[-1] if all_versions else None
        }
    
    async def reset_database(self):
        """重置数据库（删除所有表）"""
        logger.warning("Resetting database - all data will be lost")
        
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.drop_all)
        
        # 清理迁移记录
        async with get_db_session() as session:
            await session.execute("DELETE FROM migration_history")
            await session.commit()
        
        self.applied_migrations.clear()
        logger.info("Database reset completed")


# 全局迁移运行器
migration_runner = MigrationRunner()


# 迁移基类
class Migration:
    """迁移基类"""
    
    def __init__(self, version: str, description: str):
        self.version = version
        self.description = description
    
    async def up(self):
        """执行迁移"""
        raise NotImplementedError("Migration.up() must be implemented")
    
    async def down(self):
        """回滚迁移"""
        raise NotImplementedError("Migration.down() must be implemented")


# 装饰器
def migration(version: str, description: str):
    """迁移装饰器"""
    def decorator(cls):
        if not issubclass(cls, Migration):
            cls = type(cls.__name__, (Migration, cls), {})
        
        instance = cls(version, description)
        migration_runner.register_migration(version, instance)
        return cls
    
    return decorator


# CLI命令
async def migrate_command(target_version: Optional[str] = None):
    """迁移命令"""
    try:
        await migration_runner.migrate(target_version)
        status = await migration_runner.get_status()
        
        if target_version:
            print(f"Migrated to version {target_version}")
        else:
            print("Migrated to latest version")
        
        print(f"Applied {status['applied_migrations']} migrations")
        
    except Exception as e:
        print(f"Migration failed: {e}")
        return False
    
    return True


async def rollback_command(target_version: str):
    """回滚命令"""
    try:
        await migration_runner.rollback(target_version)
        print(f"Rolled back to version {target_version}")
        
    except Exception as e:
        print(f"Rollback failed: {e}")
        return False
    
    return True


async def status_command():
    """状态命令"""
    try:
        status = await migration_runner.get_status()
        
        print(f"Total migrations: {status['total_migrations']}")
        print(f"Applied migrations: {status['applied_migrations']}")
        print(f"Pending migrations: {status['pending_migrations']}")
        print(f"Current version: {status['current_version']}")
        print(f"Latest version: {status['latest_version']}")
        
        if status['pending_migrations']:
            print("\nPending migrations:")
            for version in status['pending_migrations']:
                print(f"  - {version}")
        
    except Exception as e:
        print(f"Status command failed: {e}")
        return False
    
    return True


async def reset_command():
    """重置命令"""
    try:
        await migration_runner.reset_database()
        print("Database reset successfully")
        
    except Exception as e:
        print(f"Reset failed: {e}")
        return False
    
    return True


# 主函数
async def main():
    """主函数"""
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python -m migrations <command> [options]")
        print("Commands:")
        print("  migrate [version]  - Apply migrations (to latest or specific version)")
        print("  rollback <version> - Rollback to specific version")
        print("  status            - Show migration status")
        print("  reset             - Reset database (CAUTION: deletes all data)")
        return
    
    command = sys.argv[1]
    
    if command == "migrate":
        target_version = sys.argv[2] if len(sys.argv) > 2 else None
        await migrate_command(target_version)
    elif command == "rollback":
        if len(sys.argv) < 3:
            print("Usage: python -m migrations rollback <version>")
            return
        target_version = sys.argv[2]
        await rollback_command(target_version)
    elif command == "status":
        await status_command()
    elif command == "reset":
        confirm = input("Are you sure you want to reset the database? This will delete all data. (yes/no): ")
        if confirm.lower() == 'yes':
            await reset_command()
        else:
            print("Reset cancelled")
    else:
        print(f"Unknown command: {command}")


if __name__ == "__main__":
    asyncio.run(main())
